package uebung07_w;

import java.awt.Image;
import java.awt.geom.AffineTransform;
import java.awt.image.BufferedImage;
import java.util.LinkedList;
import java.util.Vector;

import javax.imageio.ImageIO;
import javax.swing.JFrame;
import javax.swing.JPanel;

import org.ejml.simple.SimpleMatrix;

import tools.ShowImage;

public class DisparityMap {
	
	public static SimpleMatrix calcAveragedDisparityEstimationWithPruning(SimpleMatrix l, SimpleMatrix r, int w, int h, int win, int d_max, double threshold, int[] pos) {
		
		/**
		 * Matrix r muss mit passendem win (bzw. d_max) gepadded sein
		 * win muss ungerade sein >= 1
		 * w und h sind weite und hoehe der linken matrix
		 * 
		 * 1.	For each pixel (x,y) in the left image,
		 * 		search for the corresponding pixel in the right image:
		 * 		I.		FOR d' = -d_max TO +d_max DO
		 * 					IF |f_l(x,y)-f_r(x+d',y)|<threshold
		 * 		II.		Calculate window costs W_c(x,y,d') == min(SSD)
		 * 		III.	Determine average disparity d_avg
		 * 				END
		 * 2.	Repeat 1. to calculate disparities of all pixels in the left image
		 * 
		 */
		
		SimpleMatrix DispMap = new SimpleMatrix(h,w);
		
		int rpos = pos[0], cpos = pos[1];
		win = (win-1)/2;
		
		for (int row=win;row<h-win;row++) { for (int col=d_max+win;col<w-d_max-win;col++) {
//		for (int row=win;row<h+win;row++) { for (int col=win;col<w+win;col++) {
			Vector<Double> wincosts = new Vector<Double>();
			for (int d = -d_max;d<=d_max;d++) {
				if ( Math.abs(l.get(row, col)-r.get(row+rpos, col+d+cpos)) < threshold) {
					double ssd = 0, tmp; // ssd = sum of squared differences
					// calculate window cost
					for (int wrow = -win;wrow<=win;wrow++) { for (int wcol=-win;wcol<=win;wcol++) {
						double tmp1 = l.get(row+wrow, col+wcol);
						double tmp2 = r.get(row+wrow+rpos, col+wcol+d+cpos);
						tmp = tmp1-tmp2; 
						ssd += tmp*tmp;
					}}
					wincosts.add(ssd);
				}
			} // END Window

			// determine average disparity d_avg and write into DisparityMap-Matrix
			double sum = 0;
			
			for (int i = 0; i<wincosts.size();i++) { sum += wincosts.get(i); }
			
			double d_avg =  sum/wincosts.size();
			DispMap.set(row, col, d_avg);
		}}

		
		return DispMap;
	}

	public static SimpleMatrix padding(SimpleMatrix M, int win, int d_max, int[] pos){
		int rowpadd = win-1;
		int rows = M.numRows()+rowpadd;
		
		int colpadd = d_max>rowpadd/2 ? d_max*2 : rowpadd;
		int cols = M.numCols()+colpadd;
		
		SimpleMatrix paddedM = new SimpleMatrix(rows, cols);
		int rpos = rowpadd/2, cpos = colpadd/2;
		paddedM.insertIntoThis(rpos, cpos, M);
		
		pos[0]=rpos+1; pos[1]=cpos+1;
		
		return paddedM;
	}
	
	
	public static void main(String[] args) throws Exception {
		BufferedImage left = ImageIO.read(DisparityMap.class.getClassLoader().getResourceAsStream("res/im2.png"));
		BufferedImage right = ImageIO.read(DisparityMap.class.getClassLoader().getResourceAsStream("res/im6.png"));
		
		int w = left.getWidth(), h = left.getHeight();
		
		int[] a = left.getRGB(0, 0, w, h, new int[w*h], 0, w);
		int[] b = right.getRGB(0, 0, w, h, new int[w*h], 0, w);
		
		int[] aa = tools.ColorConversions.getFastLuminance(a, new int[w*h]);
		int[] bb = tools.ColorConversions.getFastLuminance(b, new int[w*h]);
		
		double[] aaa = tools.Image.getScalarLuminance(aa, new double[w*h]);
		double[] bbb = tools.Image.getScalarLuminance(bb, new double[w*h]);
		
		SimpleMatrix l = new SimpleMatrix(h, w);
		for (int i=0;i<aaa.length;i++) l.set(i, aaa[i]);
		SimpleMatrix r = new SimpleMatrix(h, w);
		for (int i=0;i<bbb.length;i++) r.set(i, bbb[i]);
		
		
		
		
		int win = 11, d_max = 2;
		double threshold = 55;
		
//		SimpleMatrix lpadded = padding(l,win,0,new int[2]);
		
		int[] pos = new int[2]; // enthaelt [row,col] die startposition in der gepaddeten matrix
		SimpleMatrix rpadded = padding(r,win,d_max,pos);
		
		SimpleMatrix DispMap = calcAveragedDisparityEstimationWithPruning(l, rpadded, w, h, win, d_max, threshold,pos);
		
		double[] preimage = new double[w*h];
		for (int i = 0; i<w*h; i++){
			preimage[i] = DispMap.get(i);
		}
		int[] pixels = tools.Image.prepare(preimage, new int[w*h]);
		for (int i = 0; i<pixels.length; i++)
			System.out.println(pixels[i]);
		int[] poxels = tools.Image.buildRGB(pixels, pixels, pixels, new int[pixels.length]);
		final BufferedImage bim = new BufferedImage(w, h, BufferedImage.TYPE_INT_ARGB);
		bim.setRGB(0, 0, w, h, poxels, 0, w);
		
		
		JFrame ff = new JFrame();
		ff.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
		ShowImage si = new ShowImage(bim);
		
		ff.add(si);
		ff.pack();
		ff.setVisible(true);
		
		
	}
	
}
